#include "mapreduce.h"

int NUM_MAPPER;                 // Mapper 个数
int NUM_REDUCER;                // Reducer 个数

char** file_names;              // 文件指针
int NUM_FILE;                   // 输入文件数

Mapper func_map;
Reducer func_reduce;
Partitioner func_partition;

typedef struct v_node {
    char *value;
    struct v_node *next;
} v_node;

typedef struct k_node {
    char *key;
    v_node *head, *curr;               // 连接 key 对应的 value 列表, curr 用于标记访问位置
    struct k_node *next;
} k_node;

typedef struct k_entry {
    k_node *head;
    pthread_mutex_t lk;     
} k_entry;

typedef struct MR_t {
    k_entry *map;               // 指向 k_entry 型数组
    pthread_mutex_t lk;         // 计数器锁
    int num_key;                // 不同 key 的个数
    k_node **sorted;            // 按字典序从小到大排列
    int curr_visted;            // 标记访问的位置
} MR_t;

MR_t *table;

char* get_next(char *key, int partition_number) {
    k_node** arr = table[partition_number].sorted;
    
    int curr = table[partition_number].curr_visted;
    if (curr == table[partition_number].num_key) return NULL;   // 已经访问完了此 MR 的所有 key

    if (strcmp(arr[curr]->key, key) == 0) {
        if (arr[curr]->curr == NULL) {
            table[partition_number].curr_visted++;
            return NULL;
        }
            
        char* value = arr[curr]->curr->value;
        arr[curr]->curr = arr[curr]->curr->next;
        return value;
    }
    else {
        table[partition_number].curr_visted++;
    }

    return NULL;
}

// 按字典序从小到大排序
void quick_sort(k_node** arr, int l, int r) {
    if (l >= r) return;
    int i = l - 1, j = r + 1, x = (l + r) >> 1;
    
    while (i < j) {
        while (strcmp(arr[++i]->key, arr[x]->key) < 0);
        while (strcmp(arr[--j]->key, arr[x]->key) > 0);
        if (i < j) {
            k_node* tmp = arr[i];
            arr[i] = arr[j];
            arr[j] = tmp;
        }
    }
    quick_sort(arr, l, j);
    quick_sort(arr, j + 1, r);
}

// 释放内存空间
void MR_Free(int partition_number) {
    for (int i = 0; i < table[partition_number].num_key; i++) {
        v_node* v_cur = table[partition_number].sorted[i]->head;
        while (v_cur) {
            v_node* v_tmp = v_cur;
            v_cur = v_cur->next;
            free(v_tmp->value);
            free(v_tmp);
        }
        free(table[partition_number].sorted[i]->key);
        free(table[partition_number].sorted[i]);
    }
    free(table[partition_number].sorted);
}

void* Reduce_Job(void *arg) {
    int partition_number = *(int*)arg;
    free(arg);

    int total = table[partition_number].num_key;

    if (total == 0) return NULL;
    
    table[partition_number].sorted = calloc(total, sizeof(k_node*));

    int count = 0;
    for (int i = 0; i < NUM_MAPPER; i++) {
        k_node *cur = table[partition_number].map[i].head;
        while (cur) {
            table[partition_number].sorted[count++] = cur;
            cur = cur->next;
        }
    }

    quick_sort(table[partition_number].sorted, 0, count - 1);
    table[partition_number].curr_visted = 0;

    for (int i = 0; i < count; i++) {
        char* key = table[partition_number].sorted[i]->key;
        table[partition_number].sorted[i]->curr = table[partition_number].sorted[i]->head;
        (*func_reduce)(key, get_next, partition_number);
    }

    MR_Free(partition_number);

    return NULL;
}

unsigned long MR_DefaultHashPartition(char *key, int num_partitions) {
    unsigned long hash = 5381;
    char c = '\0';
    while ((c == *key++) != '\0')
        hash = hash * 33 + c;
    return hash % num_partitions;
}

// Mapper 的任务
void* Map_Job(void *arg) {
    int idx = *(int*)arg;
    free(arg);
    
    while (idx < NUM_FILE) {
        (*func_map)(file_names[idx]);
        idx += NUM_MAPPER;
    }

    return NULL;
}

// 存储 key/value 对
void MR_Emit(char *key, char *value) {
    // 检查合法性
    int n = strlen(key), m = strlen(value);
    if (n == 0 || m == 0) return;

    unsigned long partition_number = (*func_partition)(key, NUM_REDUCER);
    unsigned long map_number = MR_DefaultHashPartition(key, NUM_MAPPER);

    pthread_mutex_lock(&table[partition_number].map[map_number].lk);

    k_node* cur = table[partition_number].map[map_number].head;

    while (cur) {
        if (strcmp(cur->key, key) == 0)
            break;
        cur = cur->next;
    }

    // 创建 value 结点
    v_node* new_v = malloc(sizeof(v_node));
    new_v->value = calloc(m + 1, sizeof(char));
    strcpy(new_v->value, value);
    new_v->next = NULL;

    if (cur == NULL) {  // key 不存在
        k_node* new_k = malloc(sizeof(k_node));
        new_k->head = new_v;
        new_k->key = calloc(n + 1, sizeof(char));
        strcpy(new_k->key, key);
        new_k->next = table[partition_number].map[map_number].head;
        table[partition_number].map[map_number].head = new_k;

        // 计数器更新
        pthread_mutex_lock(&table[partition_number].lk);
        table[partition_number].num_key++;
        pthread_mutex_unlock(&table[partition_number].lk);
    }
    else {
        new_v->next = cur->head;
        cur->head = new_v;
    }

    pthread_mutex_unlock(&table[partition_number].map[map_number].lk);
}


void MR_init(int argc, char *argv[], Mapper map, int num_mappers, Reducer reduce, int num_reducers, Partitioner patition) {
    NUM_FILE = argc - 1;
    file_names = argv + 1;

    func_map = map;
    NUM_MAPPER = num_mappers;

    func_reduce = reduce;
    NUM_REDUCER = num_reducers;

    func_partition = patition;

    table = malloc(num_reducers * sizeof(MR_t));

    for (int i = 0; i < num_reducers; i++) {
        table[i].map = calloc(num_mappers, sizeof(k_entry));
        table[i].num_key = 0;
        table[i].sorted = NULL;
        pthread_mutex_init(&table[i].lk, NULL);

        // 初始化互斥锁
        for (int j = 0; j < num_mappers; j++) {
            table[i].map[j].head = NULL;
            pthread_mutex_init(&table[i].map[j].lk, NULL);
        }
    }
}

void MR_Run(int argc, char *argv[], Mapper map, int num_mappers, Reducer reduce, int num_reducers, Partitioner patition) {
    // 初始化
    MR_init(argc, argv, map, num_mappers, reduce, num_reducers, patition);

    pthread_t mappers[num_mappers];

    for (int i = 0; i < num_mappers; i++) {
        int* file_idx = malloc(sizeof(int));
        *file_idx = i;
        pthread_create(&mappers[i], NULL, Map_Job, (void*)file_idx);
    }

    // 等待所有 Map 线程结束
    for (int i = 0; i < num_mappers; i++) {
        pthread_join(mappers[i], NULL);
    }

    pthread_t reducers[num_reducers];

    for (int i = 0; i < num_reducers; i++) {
        int* partition_number = malloc(sizeof(int));
        *partition_number = i;
        pthread_create(&reducers[i], NULL, Reduce_Job, (void*)partition_number);
    }

    // 等待所有 Reduce 线程结束
    for (int i = 0; i < num_reducers; i++) {
        pthread_join(reducers[i], NULL);
    }

    for (int i = 0; i < num_reducers; i++) {
        free(table[i].map);
    }

    free(table);
}

